from sklearn.datasets import load_iris
from sklearn.naive_bayes import GaussianNB
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report
import matplotlib.pyplot as plt
import numpy as np

# 加载鸢尾花数据集
iris = load_iris()
X = iris.data  # 特征：花萼长度、花萼宽度、花瓣长度、花瓣宽度
y = iris.target  # 标签：0-Setosa, 1-Versicolour, 2-Virginica
feature_names = iris.feature_names
target_names = iris.target_names

print("数据集信息:")
print(f"特征: {feature_names}")
print(f"类别: {list(target_names)}")
print(f"数据形状: {X.shape}")

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# 创建并训练高斯朴素贝叶斯模型（适用于连续数据）
model = GaussianNB()
model.fit(X_train, y_train)

# 预测
y_pred = model.predict(X_test)

# 评估模型
accuracy = accuracy_score(y_test, y_pred)
print(f"\n准确率: {accuracy:.2f}")
print(f"\n分类报告:")
print(classification_report(y_test, y_pred, target_names=target_names))

# 可视化部分结果（前两个特征）
plt.figure(figsize=(12, 5))

# 真实标签
plt.subplot(1, 2, 1)
for i, color in zip([0, 1, 2], ['red', 'blue', 'green']):
    plt.scatter(X_test[y_test == i, 0], X_test[y_test == i, 1], 
               c=color, label=target_names[i], alpha=0.7)
plt.xlabel(feature_names[0])
plt.ylabel(feature_names[1])
plt.title('True Labels')
plt.legend()

# 预测标签
plt.subplot(1, 2, 2)
for i, color in zip([0, 1, 2], ['red', 'blue', 'green']):
    plt.scatter(X_test[y_pred == i, 0], X_test[y_pred == i, 1], 
               c=color, label=target_names[i], alpha=0.7)
plt.xlabel(feature_names[0])
plt.ylabel(feature_names[1])
plt.title('Predicted Labels')
plt.legend()

plt.tight_layout()
plt.show()

# 测试新样本
new_samples = np.array([
    [5.1, 3.5, 1.4, 0.2],  # 类似Setosa
    [6.0, 2.7, 5.1, 1.6],  # 类似Virginica
    [5.5, 2.4, 3.8, 1.1]   # 类似Versicolour
])

new_predictions = model.predict(new_samples)
new_probabilities = model.predict_proba(new_samples)

print("\n新样本预测:")
for i, (sample, pred, prob) in enumerate(zip(new_samples, new_predictions, new_probabilities)):
    print(f"样本 {i+1}: {sample}")
    print(f"预测类别: {target_names[pred]}")
    print(f"各类别概率: {dict(zip(target_names, prob))}")
    print()